#!/usr/bin/python
import sys, os, re, time
################################################################################
################################################################################
class Params:
   trial   = False
   verbose = True
   rebuild = False
   mainmem = 96*1024*1024*1024
   memfact = 20
   threads = [ 24, 16, 8, 4, 2, 1 ]
   def __init__( self, mtxpath ):
      
      assert os.path.isdir( mtxpath )
      self.mtxpath = mtxpath
      dagHomeEnvName = 'DAGMTX_HOME_DIR'
      
      if not dagHomeEnvName in os.environ:
         print 'Before using ' + sys.argv[0] + ', you need to set the environment variable "' + dagHomeEnvName + '"'
         print 'to point to the home directory for this project.'
         sys.exit( -1 )
      
      self.dagHomeDir = os.environ[ dagHomeEnvName ]      
      self.binHome = os.path.join( self.dagHomeDir, 'bin' )
      self.unzBin  = os.path.join( self.binHome,    'unz' )
   

################################################################################
################################################################################
def ReadFileLines( file ):
   # opens a file
   # returns the contents as a list of lines in the file
   # trailing new-line chars are removed
   FH = open( file, 'r' )
   lines = [ line.rstrip() for line in FH.readlines() ]
   FH.close()
   return lines

def AddRedirectionToCmd( cmd, fn ):
  return " ".join([ cmd, "1>", fn, "2>&1" ])

def PrintTimeStamp():
   theTime = time.localtime()
   yr  = '%04d' % theTime.tm_year
   mon = '%02d' % theTime.tm_mon
   day = '%02d' % theTime.tm_mday
   hr  = '%02d' % theTime.tm_hour
   mt  = '%02d' % theTime.tm_min
   sec = '%02d' % theTime.tm_sec
   print '...', '/'.join( [yr, mon, day] ), ':'.join( [hr, mt, sec] )
   sys.stdout.flush()

def SystemWrapperPipelined( cmd, verbose, trial = False ):
   if verbose:
      PrintTimeStamp()
      print '... os.popen(', cmd, ')'
      sys.stdout.flush()
   if not trial:
      FH = os.popen( cmd )
      output = [ line.rstrip() for line in FH.readlines() ]
      FH.close()
      return output
   return []

def SystemWrapper( cmd, verbose, trial = False ):
   dt = 0
   if verbose:
      PrintTimeStamp()
      print '... os.system(', cmd, ')'
      sys.stdout.flush()
   if trial:
      status = 0
   else:
      t0 = time.time()
      status = os.system( cmd )
      t1 = time.time()
      dt = t1 - t0
   assert status == 0
   if verbose:
      print '...', dt, 'seconds elapsed during system call.'
   return dt

def RemoveWrapper( fn, verbose, trial = False ):
	if verbose:
		PrintTimeStamp()
		print '... os.remove(', fn, ')'
		sys.stdout.flush()
	if not trial:
		os.remove( fn )

################################################################################
################################################################################
def HeadNLines( fn, n ):
   cmd = ' '.join([ 'head', '-%d' % n, fn ])
   return SystemWrapperPipelined( cmd, False, False )

def IsRealOrIntMtx( mtxFile ):
	line = HeadNLines( mtxFile, 1 )[0].rstrip().lstrip( '%' )
	toks = line.split()
	if len( toks ) != 5:          return False
	if toks[0] != 'MatrixMarket': return False
	if toks[1] != 'matrix':       return False
	if toks[2] != 'coordinate':   return False
	if toks[3] not in [    'real',   'integer' ]: return False
	if toks[4] not in [ 'general', 'symmetric' ]: return False
	return True

def MtxDims( mtxFile ):
   lines = HeadNLines( mtxFile, 200 )
   for line in lines:
      if len( line ):
         if line[0] != '%': break
   toks = line.split()
   assert len( toks ) == 3
   m   = int( toks[0] )
   n   = int( toks[1] )
   nnz = int( toks[2] )
   return ( m, n, nnz )

def CSRBytes( m, n, nnz ):
   nnzValBytes = 8 * nnz
   colIdxBytes = 4 * nnz
   rowPtrBytes = 4 * m + 4
   totalBytes = nnzValBytes + colIdxBytes + rowPtrBytes
   return totalBytes

################################################################################
################################################################################
def ForkLaunch( cmds, logs, maxFork = 1, verbose = True, trial = False ):
   # sanity:
   assert len( cmds ) == len( logs )
   #
   # setup:
   cmdslogs = zip( cmds, logs )
   numLaunched = 0
   #
   # fork & launch:
   for ( cmd, log ) in cmdslogs:
      pid = os.fork()
      cmd = AddRedirectionToCmd( cmd, log )
      if pid == 0:
         SystemWrapper( cmd, verbose = verbose, trial = trial )
         sys.exit( 0 )
      else:
         numLaunched += 1
         if numLaunched >= maxFork:
            os.waitpid( -1, 0 )
   while True:
      try:    os.waitpid( -1, 0 )
      except: break

################################################################################
################################################################################
def DagFileExists( mtxFile, fmt, deg ):
   parts = 1
   ( path, fn ) = os.path.split( mtxFile )
   toks = fn.split( '.' )
   toks.pop()
   rawf = '.'.join( toks )
   rawdag = '%s.%s.deg%02d.parts%02d.dag' % ( rawf, fmt.lower(), deg, parts )
   dagfn  = os.path.join( path, rawdag )
   if os.path.isfile( dagfn ): return True
   return False

def MemLimit( params, numThreads ):
   if numThreads == 1: return 1024*1024*1024*1024
   return params.mainmem / numThreads / params.memfact

def GenCmdLogTup( params, mtxFile ):
   parts = 1
   ( path, fn ) = os.path.split( params.unzBin )
   toks  = fn.split( '.' )
   rawc  = toks[0]
   ( path, fn ) = os.path.split( mtxFile )
   toks  = fn.split( '.' )
   rawf  = toks[0]
   cmd   = ' '.join( [ params.unzBin, mtxFile ] )
   log   = '.'.join( [ rawc, rawf, 'log' ] )
   log   = os.path.join( path, log )
   return ( cmd, log )

def GenerateCmdAndLogList( params, inMtxFileTups, numThreads ):
   mtxFileTups = [ x for x in inMtxFileTups ]
   cmdLogList = []
   memLimit = MemLimit( params, numThreads )
   while len( mtxFileTups ):
      tup = mtxFileTups.pop()
      ( csrBytes, mtxFile, m, n, nnz ) = tup
      if csrBytes < memLimit:
         ( cmd, log ) = GenCmdLogTup( params, mtxFile )
         if not os.path.isfile( log ):
            cmdLogList.append( ( cmd, log ) )
      else:
         mtxFileTups.append( tup )
         break
   return cmdLogList

def CountUNZs( params, mtxFileTups ):
   for numThreads in params.threads:
      cmdslogs = GenerateCmdAndLogList( params, mtxFileTups, numThreads )
      if cmdslogs == []: continue
      ( cmds, logs ) = zip( *cmdslogs )
      ForkLaunch( cmds, logs, maxFork = numThreads, verbose = params.verbose, trial = params.trial )

################################################################################
################################################################################
def SortAndSizeMtxs( mtxFiles ):
   tups = []
   for mtxFile in mtxFiles:
      ( m, n, nnz ) = MtxDims( mtxFile )
      csrBytes = CSRBytes( m, n, nnz )
      tup = ( csrBytes, mtxFile, m, n, nnz )
      tups.append( tup )
   tups.sort()
   tups.reverse()
   return tups

def WalkAndFindMtxs( params ):
   mtxFiles = []
   for ( dirpath, dirs, files ) in os.walk( params.mtxpath ):
      for fn in files:
         toks = fn.split( '.' )
         if toks[ -1 ] == 'mtx':
            mtxFile = os.path.join( dirpath, fn )
            if IsRealOrIntMtx( mtxFile ):
               mtxFiles.append( mtxFile )
   return mtxFiles

################################################################################
################################################################################
def MtxList():
   ml = []
   ml.append( 'bayer02'          )
   ml.append( 'powersim'         )
   ml.append( 'bayer10'          )
   ml.append( 'rdist1'           )
   ml.append( 'memplus'          )
   ml.append( 'mark3jac040'      )
   ml.append( 'sparseAI'         )
   ml.append( 'shermanACb'       )
   ml.append( 'wang3'            )
   ml.append( 'wang4'            )
   ml.append( 'coater2'          )
   ml.append( 'Ill_Stokes'       )
   ml.append( 'lhr10'            )
   ml.append( 'ex35'             )
   ml.append( 'nmos3'            )
   ml.append( 'onetone2'         )
   ml.append( 'lp_cre_d'         )
   ml.append( 'chipcool0'        )
   ml.append( 'goodwin'          )
   ml.append( 'lp_osa_14'        )
   ml.append( 'g7jac100'         )
   ml.append( 'lhr17'            )
   ml.append( 'lhr17c'           )
   ml.append( 'shyy161'          )
   ml.append( 'bcircuit'         )
   ml.append( 'af23560'          )
   ml.append( 'lung2'            )
   ml.append( 'lp_osa_30'        )
   ml.append( 'hcircuit'         )
   ml.append( 'Baumann'          )
   ml.append( 'sinc18'           )
   ml.append( 'rim'              )
   ml.append( 'matrix-new_3'     )
   ml.append( 'ex11'             )
   ml.append( 'scircuit'         )
   ml.append( 'std1_Jac2'        )
   ml.append( 'Maragal_8'        )
   ml.append( 'twotone'          )
   ml.append( 'raefsky3'         )
   ml.append( 'lp_osa_60'        )
   ml.append( 'mac_econ_fwd500'  )
   ml.append( 'ns3Da'            )
   ml.append( 'av41092'          )
   ml.append( 'venkat01'         )
   ml.append( 'majorbasis'       )
   ml.append( 'crashbasis'       )
   ml.append( 'barrier2-1'       )
   ml.append( 'poisson3Db'       )
   ml.append( 'mc2depi'          )
   ml.append( 'TSOPF_RS_b300_c2' )
   ml.append( 'thermomech_dK'    )
   ml.append( 'sme3Dc'           )
   ml.append( 'stat96v3'         )
   ml.append( 'xenon2'           )
   ml.append( 'webbase-1M'       )
   ml.append( 'rajat29'          )
   ml.append( 'stormG2_1000'     )
   ml.append( 'Chebyshev4'       )
   ml.append( 'largebasis'       )
   ml.append( 'pre2'             )
   ml.append( 'ohne2'            )
   ml.append( 'Hamrle3'          )
   ml.append( 'PR02R'            )
   ml.append( 'torso1'           )
   ml.append( 'Rucci1'           )
   ml.append( 'tp-6'             )
   ml.append( 'atmosmodl'        )
   ml.append( 'TSOPF_RS_b2383'   )
   ml.append( 'circuit5M_dc'     )
   ml.append( 'rajat31'          )
   ml.append( 'cage14'           )
   ml.append( 'FullChip'         )
   ml.append( 'RM07R'            )
   ml.append( 'circuit5M'        )
   return ml

################################################################################
################################################################################
def ReadLogFile( logf ):
   nnz = 0
   unz = 0
   lines = ReadFileLines( logf )
   for line in lines:
      toks = line.split()
      if toks[0] == 'nnz':
         assert nnz == 0
         nnz = int( toks[2] )
      if toks[0] == 'unz':
         assert unz == 0
         unz = int( toks[2] )
   return ( nnz, unz )

def GatherForMatlab( mtxFileTups ):
   unzMap = {}
   for ( csrBytes, mtxFile, m, n, nnz ) in mtxFileTups:
      ( path, fn ) = os.path.split( mtxFile )
      toks = fn.split( '.' )
      rawf = '.'.join( toks[ :-1] )
      logf = '.'.join([ 'unz', rawf, 'log' ])
      logf = os.path.join( path, logf )
      ( nnz, unz ) = ReadLogFile( logf )
      unzMap[ rawf ] = ( nnz, unz )
   return unzMap

def PrintToMatlab( unzMap, mtxList ):
   vnnz = []
   vunz = []
   for mtx in mtxList:
      ( nnz, unz ) = unzMap[ mtx ]
      assert nnz
      assert unz
      assert unz <= nnz
      vnnz.append( str( nnz ) )
      vunz.append( str( unz ) )
   print 'vnnz = [ ' + ', '.join( vnnz ) + ' ];'
   print 'vunz = [ ' + ', '.join( vunz ) + ' ];'

################################################################################
################################################################################
if len( sys.argv ) < 2:
   ( progPath, progName ) = os.path.split( sys.argv[0] )
   print 'Usage:', progName, '<pathToMtxs>'
   sys.exit( -1 )

params = Params( sys.argv[1] )

mtxFiles    = WalkAndFindMtxs( params )
mtxFileTups = SortAndSizeMtxs( mtxFiles )
CountUNZs( params, mtxFileTups )
unzMap  = GatherForMatlab( mtxFileTups )
mtxList = MtxList()
PrintToMatlab( unzMap, mtxList )


